import sys

sys.path.insert(0, '../src')

import torch
from inets import FrozenInTime
from inets import Tokenizer


if __name__ == '__main__':
    video_params = {
        "model": "SpaceTimeTransformer",
        "arch_config": "base_patch16_224",
        "num_frames": 4,
        "pretrained": True,
        "time_init": "zeros"
    }

    text_params = {
        "model": "distilbert-base-uncased",
        "pretrained": True,
        "input": "text"
    }
    net = FrozenInTime(video_params, text_params, projection='minimal', load_checkpoint="")

    video = torch.randn((2, 4, 3, 224, 224))

    ve = net.compute_video(video)
    print(ve)

    t = Tokenizer(text_params['model'], model_max_length=1e6, TOKENIZERS_PARALLELISM=False)

    text = "i say the word is not good"
    i = t(text)
    print(i)

